"""
Database manager for SkyWalking data collection system.
Provides CRUD operations with connection pooling and transaction management.
"""

import os
import json
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Any, Tuple
from contextlib import contextmanager

from sqlalchemy import create_engine, and_, or_, func, desc
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool

from .schema import (
    Base, Service, ServiceMetric, Endpoint, EndpointMetric, 
    Trace, CollectionLog, create_tables
)


class DatabaseManager:
    """
    数据库管理器类
    支持 SQLite 和 PostgreSQL，通过配置切换
    """
    
    def __init__(self, database_url: str, echo: bool = False):
        """
        初始化数据库管理器
        
        Args:
            database_url: 数据库连接 URL
                - SQLite: sqlite:///path/to/db.db
                - PostgreSQL: postgresql://user:pass@host:5432/dbname
            echo: 是否输出 SQL 日志
        """
        self.database_url = database_url
        
        # 为 SQLite 配置连接参数
        if database_url.startswith('sqlite'):
            # 确保数据库目录存在
            db_path = database_url.replace('sqlite:///', '')
            os.makedirs(os.path.dirname(db_path) if os.path.dirname(db_path) else '.', exist_ok=True)
            
            self.engine = create_engine(
                database_url,
                echo=echo,
                connect_args={'check_same_thread': False}  # SQLite specific
            )
        else:
            # PostgreSQL 使用连接池
            self.engine = create_engine(
                database_url,
                echo=echo,
                poolclass=QueuePool,
                pool_size=5,
                max_overflow=10,
                pool_pre_ping=True  # 验证连接有效性
            )
        
        self.SessionLocal = sessionmaker(bind=self.engine, autoflush=False, autocommit=False)
        
        # 创建所有表
        create_tables(self.engine)
    
    @contextmanager
    def get_session(self) -> Session:
        """
        获取数据库会话（上下文管理器）
        自动处理事务提交和回滚
        """
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
    
    # ==================== Service 操作 ====================
    
    def upsert_service(self, service_id: str, name: str, layer: str = None, group: str = None) -> Service:
        """
        插入或更新服务信息
        
        Returns:
            Service: 服务对象
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if service:
                # 更新
                service.name = name
                service.layer = layer
                service.group = group
                service.collected_at = datetime.utcnow()
            else:
                # 插入
                service = Service(
                    service_id=service_id,
                    name=name,
                    layer=layer,
                    group=group
                )
                session.add(service)
            session.flush()
            return service
    
    def get_service_by_id(self, service_id: str) -> Optional[Service]:
        """根据 SkyWalking service_id 获取服务"""
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if service:
                # 触发所有属性加载，避免 DetachedInstanceError
                _ = service.id, service.service_id, service.name, service.layer, service.group
            return service
    
    def get_all_services(self) -> List[Service]:
        """获取所有服务"""
        with self.get_session() as session:
            services = session.query(Service).order_by(Service.name).all()
            # 触发所有属性加载，避免 DetachedInstanceError
            for s in services:
                _ = s.id, s.service_id, s.name, s.layer, s.group
            return services
    
    # ==================== ServiceMetric 操作 ====================
    
    def insert_service_metrics(self, service_id: str, metrics: List[Dict[str, Any]]) -> int:
        """
        批量插入服务指标
        
        Args:
            service_id: SkyWalking service ID
            metrics: 指标列表，格式: [{'metric_name': 'xxx', 'value': 123, 'timestamp': datetime}, ...]
        
        Returns:
            int: 插入的记录数
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                raise ValueError(f"Service not found: {service_id}")
            
            metric_objects = []
            for metric in metrics:
                metric_objects.append(ServiceMetric(
                    service_id=service.id,
                    metric_name=metric['metric_name'],
                    value=metric['value'],
                    timestamp=metric['timestamp']
                ))
            
            session.bulk_save_objects(metric_objects)
            return len(metric_objects)
    
    def get_service_metrics(
        self, 
        service_id: str, 
        metric_names: List[str] = None,
        start_time: datetime = None,
        end_time: datetime = None
    ) -> List[ServiceMetric]:
        """
        获取服务指标
        
        Args:
            service_id: SkyWalking service ID
            metric_names: 指标名称列表，None 表示所有
            start_time: 开始时间
            end_time: 结束时间
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                return []
            
            query = session.query(ServiceMetric).filter(ServiceMetric.service_id == service.id)
            
            if metric_names:
                query = query.filter(ServiceMetric.metric_name.in_(metric_names))
            if start_time:
                query = query.filter(ServiceMetric.timestamp >= start_time)
            if end_time:
                query = query.filter(ServiceMetric.timestamp <= end_time)
            
            return query.order_by(ServiceMetric.timestamp).all()
    
    # ==================== Endpoint 操作 ====================
    
    def upsert_endpoint(self, service_id: str, endpoint_id: str, name: str) -> Endpoint:
        """
        插入或更新端点信息
        
        Returns:
            Endpoint: 端点对象
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                raise ValueError(f"Service not found: {service_id}")
            
            endpoint = session.query(Endpoint).filter_by(endpoint_id=endpoint_id).first()
            if endpoint:
                # 更新
                endpoint.name = name
                endpoint.collected_at = datetime.utcnow()
            else:
                # 插入
                endpoint = Endpoint(
                    service_id=service.id,
                    endpoint_id=endpoint_id,
                    name=name
                )
                session.add(endpoint)
            session.flush()
            return endpoint
    
    def get_endpoints_by_service(self, service_id: str) -> List[Endpoint]:
        """获取服务的所有端点"""
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                return []
            endpoints = session.query(Endpoint).filter_by(service_id=service.id).order_by(Endpoint.name).all()
            # 触发所有属性加载，避免 DetachedInstanceError
            for e in endpoints:
                _ = e.id, e.endpoint_id, e.name, e.service_id
            return endpoints
    
    # ==================== EndpointMetric 操作 ====================
    
    def insert_endpoint_metrics(self, endpoint_id: str, metrics: List[Dict[str, Any]]) -> int:
        """
        批量插入端点指标
        
        Args:
            endpoint_id: SkyWalking endpoint ID
            metrics: 指标列表
        
        Returns:
            int: 插入的记录数
        """
        with self.get_session() as session:
            endpoint = session.query(Endpoint).filter_by(endpoint_id=endpoint_id).first()
            if not endpoint:
                raise ValueError(f"Endpoint not found: {endpoint_id}")
            
            metric_objects = []
            for metric in metrics:
                metric_objects.append(EndpointMetric(
                    endpoint_id=endpoint.id,
                    metric_name=metric['metric_name'],
                    value=metric['value'],
                    timestamp=metric['timestamp']
                ))
            
            session.bulk_save_objects(metric_objects)
            return len(metric_objects)
    
    def get_endpoint_metrics(
        self,
        endpoint_id: str,
        metric_names: List[str] = None,
        start_time: datetime = None,
        end_time: datetime = None
    ) -> List[EndpointMetric]:
        """获取端点指标"""
        with self.get_session() as session:
            endpoint = session.query(Endpoint).filter_by(endpoint_id=endpoint_id).first()
            if not endpoint:
                return []
            
            query = session.query(EndpointMetric).filter(EndpointMetric.endpoint_id == endpoint.id)
            
            if metric_names:
                query = query.filter(EndpointMetric.metric_name.in_(metric_names))
            if start_time:
                query = query.filter(EndpointMetric.timestamp >= start_time)
            if end_time:
                query = query.filter(EndpointMetric.timestamp <= end_time)
            
            return query.order_by(EndpointMetric.timestamp).all()
    
    # ==================== Trace 操作 ====================
    
    def insert_traces(self, service_id: str, traces: List[Dict[str, Any]]) -> int:
        """
        批量插入 Trace 记录
        
        Args:
            service_id: SkyWalking service ID
            traces: Trace 列表，格式: [{'trace_id': 'xxx', 'duration': 123, ...}, ...]
        
        Returns:
            int: 插入的记录数
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                raise ValueError(f"Service not found: {service_id}")
            
            trace_objects = []
            for trace in traces:
                # 检查是否已存在（使用 segment_id，因为它是唯一的）
                segment_id = trace.get('segment_id')
                if not segment_id:
                    continue  # 跳过没有 segment_id 的记录
                
                existing = session.query(Trace).filter_by(segment_id=segment_id).first()
                if existing:
                    continue  # 跳过重复
                
                trace_objects.append(Trace(
                    service_id=service.id,
                    trace_id=trace['trace_id'],
                    segment_id=segment_id,
                    endpoint_names=json.dumps(trace.get('endpoint_names', [])),
                    duration=trace['duration'],
                    start_time=trace['start_time'],
                    is_error=trace.get('is_error', False)
                ))
            
            session.bulk_save_objects(trace_objects)
            return len(trace_objects)
    
    def get_traces(
        self,
        service_id: str = None,
        start_time: datetime = None,
        end_time: datetime = None,
        is_error: bool = None,
        limit: int = 100
    ) -> List[Trace]:
        """
        获取 Trace 记录
        
        Args:
            service_id: SkyWalking service ID
            start_time: 开始时间
            end_time: 结束时间
            is_error: 是否只查询错误 Trace
            limit: 最大返回数量
        """
        with self.get_session() as session:
            query = session.query(Trace)
            
            if service_id:
                service = session.query(Service).filter_by(service_id=service_id).first()
                if service:
                    query = query.filter(Trace.service_id == service.id)
                else:
                    return []
            
            if start_time:
                query = query.filter(Trace.start_time >= start_time)
            if end_time:
                query = query.filter(Trace.start_time <= end_time)
            if is_error is not None:
                query = query.filter(Trace.is_error == is_error)
            
            return query.order_by(desc(Trace.start_time)).limit(limit).all()
    
    # ==================== CollectionLog 操作 ====================
    
    def insert_collection_log(
        self,
        status: str,
        message: str = None,
        services_collected: int = 0,
        endpoints_collected: int = 0,
        traces_collected: int = 0,
        metrics_collected: int = 0,
        error_details: Dict = None
    ) -> CollectionLog:
        """插入采集日志"""
        with self.get_session() as session:
            log = CollectionLog(
                status=status,
                message=message,
                services_collected=services_collected,
                endpoints_collected=endpoints_collected,
                traces_collected=traces_collected,
                metrics_collected=metrics_collected,
                error_details=json.dumps(error_details) if error_details else None
            )
            session.add(log)
            session.flush()
            return log
    
    def get_collection_logs(
        self,
        start_time: datetime = None,
        end_time: datetime = None,
        status: str = None,
        limit: int = 100
    ) -> List[CollectionLog]:
        """获取采集日志"""
        with self.get_session() as session:
            query = session.query(CollectionLog)
            
            if start_time:
                query = query.filter(CollectionLog.collected_at >= start_time)
            if end_time:
                query = query.filter(CollectionLog.collected_at <= end_time)
            if status:
                query = query.filter(CollectionLog.status == status)
            
            return query.order_by(desc(CollectionLog.collected_at)).limit(limit).all()
    
    def get_latest_collection_log(self) -> Optional[CollectionLog]:
        """获取最新的采集日志"""
        with self.get_session() as session:
            log = session.query(CollectionLog).order_by(desc(CollectionLog.collected_at)).first()
            if log:
                # 触发所有属性加载，避免 DetachedInstanceError
                _ = (log.id, log.status, log.message, log.services_collected,
                     log.endpoints_collected, log.traces_collected, log.metrics_collected,
                     log.error_details, log.collected_at)
            return log
    
    # ==================== 数据清理 ====================
    
    def cleanup_old_data(self, days: int = 90) -> Dict[str, int]:
        """
        清理超过指定天数的旧数据
        
        Args:
            days: 保留天数
        
        Returns:
            Dict: 删除的记录数统计
        """
        cutoff_date = datetime.utcnow() - timedelta(days=days)
        
        with self.get_session() as session:
            # 删除旧的 Trace
            traces_deleted = session.query(Trace).filter(Trace.collected_at < cutoff_date).delete()
            
            # 删除旧的指标
            service_metrics_deleted = session.query(ServiceMetric).filter(
                ServiceMetric.collected_at < cutoff_date
            ).delete()
            
            endpoint_metrics_deleted = session.query(EndpointMetric).filter(
                EndpointMetric.collected_at < cutoff_date
            ).delete()
            
            # 删除旧的日志
            logs_deleted = session.query(CollectionLog).filter(
                CollectionLog.collected_at < cutoff_date
            ).delete()
            
            return {
                'traces': traces_deleted,
                'service_metrics': service_metrics_deleted,
                'endpoint_metrics': endpoint_metrics_deleted,
                'collection_logs': logs_deleted
            }
    
    # ==================== 统计查询 ====================
    
    def get_service_statistics(
        self,
        service_id: str,
        start_time: datetime,
        end_time: datetime
    ) -> Dict[str, Any]:
        """
        获取服务统计信息
        
        Returns:
            Dict: 包含各类统计数据
        """
        with self.get_session() as session:
            service = session.query(Service).filter_by(service_id=service_id).first()
            if not service:
                return {}
            
            # Trace 统计
            total_traces = session.query(func.count(Trace.id)).filter(
                and_(
                    Trace.service_id == service.id,
                    Trace.start_time >= start_time,
                    Trace.start_time <= end_time
                )
            ).scalar()
            
            error_traces = session.query(func.count(Trace.id)).filter(
                and_(
                    Trace.service_id == service.id,
                    Trace.start_time >= start_time,
                    Trace.start_time <= end_time,
                    Trace.is_error == True
                )
            ).scalar()
            
            avg_duration = session.query(func.avg(Trace.duration)).filter(
                and_(
                    Trace.service_id == service.id,
                    Trace.start_time >= start_time,
                    Trace.start_time <= end_time
                )
            ).scalar()
            
            return {
                'total_traces': total_traces or 0,
                'error_traces': error_traces or 0,
                'success_rate': ((total_traces - error_traces) / total_traces * 100) if total_traces else 0,
                'avg_duration': float(avg_duration) if avg_duration else 0
            }
    
    def close(self):
        """关闭数据库连接"""
        self.engine.dispose()

